import torch
import torch.nn as nn

from vehicle_reid_pytorch.loss.triplet_loss import normalize, euclidean_dist, hard_example_mining
from ..tools.math_tools import clck_dist

class ParsingTripletLoss:
    def __init__(self, margin=None):
        self.margin = margin
        if margin is not None:
            self.ranking_loss = nn.MarginRankingLoss(margin=margin)
        else:
            self.ranking_loss = nn.SoftMarginLoss()

    def __call__(self, local_feat, vis_score, target, normalize_feature=False):
        """

        :param torch.Tensor local_feature: (B, C, N)
        :param torch.Tensor visibility_score: (B, N)
        :param torch.Tensor target: (B)
        :return:
        """
        B, C, _ = local_feat.shape
        if normalize_feature:
            local_feat = normalize(local_feat, 1)

        dist_mat = clck_dist(local_feat, local_feat,
                             vis_score, vis_score)

        dist_ap, dist_an = hard_example_mining(dist_mat, target)
        y = dist_an.new().resize_as_(dist_an).fill_(1)

        if self.margin is not None:
            loss = self.ranking_loss(dist_an, dist_ap, y)
        else:
            loss = self.ranking_loss(dist_an - dist_ap, y)

        return loss, dist_ap, dist_an